#!/usr/bin/env python
# Python Process Manager (2011r01);
# (c) 2011 by James King;
# Available under Microsoft Reciprocal License (Ms-RL)
######################################################
# By using/downloading/etc this software, you agree to the following:
"""
Microsoft Reciprocal License (Ms-RL)

This license governs use of the accompanying software. If you use the
software, you accept this license. If you do not accept the license,
do not use the software.

1. Definitions
The terms "reproduce," "reproduction," "derivative works," and
"distribution" have the same meaning here as under U.S. copyright law.

A "contribution" is the original software, or any additions or changes
to the software.

A "contributor" is any person that distributes its contribution under
this license.

"Licensed patents" are a contributor's patent claims that read directly
on its contribution.

2. Grant of Rights
(A) Copyright Grant- Subject to the terms of this license, including
the license conditions and limitations in section 3, each contributor
grants you a non-exclusive, worldwide, royalty-free copyright license
to reproduce its contribution, prepare derivative works of its
contribution, and distribute its contribution or any derivative works
that you create.

(B) Patent Grant- Subject to the terms of this license, including the
license conditions and limitations in section 3, each contributor grants
you a non-exclusive, worldwide, royalty-free license under its licensed
patents to make, have made, use, sell, offer for sale, import, and/or
otherwise dispose of its contribution in the software or derivative works
of the contribution in the software.

3. Conditions and Limitations
(A) Reciprocal Grants- For any file you distribute that contains code
from the software (in source code or binary format), you must provide
recipients the source code to that file along with a copy of this license,
which license will govern that file. You may license other files that are
entirely your own work and do not contain code from the software under any
terms you choose.

(B) No Trademark License- This license does not grant you rights to use
any contributors' name, logo, or trademarks.

(C) If you bring a patent claim against any contributor over patents that
you claim are infringed by the software, your patent license from such
contributor to the software ends automatically.

(D) If you distribute any portion of the software, you must retain all
copyright, patent, trademark, and attribution notices that are present
in the software.

(E) If you distribute any portion of the software in source code form, you
may do so only under this license by including a complete copy of this
license with your distribution. If you distribute any portion of the
software in compiled or object code form, you may only do so under a
license that complies with this license.

(F) The software is licensed "as-is." You bear the risk of using it. The
contributors give no express warranties, guarantees or conditions. You may
have additional consumer rights under your local laws which this license
cannot change. To the extent permitted under your local laws, the
contributors exclude the implied warranties of merchantability, fitness
for a particular purpose and non-infringement. 
"""

import commands;
import sys;

def processLine(line):
    processed = [];
    result = line.split(" ");
    for piece in result:
        if len(processed) < 8:
            if piece is not "":
                processed.insert(len(processed), piece);
        else:
            processed[len(processed) - 1]+= " " + piece;
    return processed;

def match(needle, haystack, args = False):
    needle = needle.replace("?", "*");
    if needle == "*": return True;

    if not args:
        haystack = haystack.split(" ")[0];
        
    if needle[0] != "/":
        haystack = haystack.split("/")[-1];

    needle = needle.split("*");
    
    if len(needle) > 2:
        if (len(needle) == 3) and (needle[0] == needle[2] == ''):
            # Case *x*
            return needle[1] in haystack;
        raise ValueError("Invalid matching string!");
    elif len(needle) > 1:
        # Case: * exists
        if len(needle) == 2:
            if (needle[0] != '') and (needle[1] != ''):
               # Case: x*y
               needle = needle[0] + needle[1];
               haystack = haystack[:len(needle[0])] + haystack[-len(needle[1]):];
            elif (needle[0] != ''):
                # Case: x*
                needle = needle[0];
                haystack = haystack[:len(needle)];
            else:
                # Case *x
                needle = needle[1];
                haystack = haystack[-len(needle):];
    else:
        needle = needle[0];
    
    if needle[-1] == "*":
        needle = needle[:-1];
        haystack = haystack[:len(needle)];

        
    #print "Needle: " + str(needle) + "; Haystack: " + str(haystack) + ";";
    if needle == haystack:
        return True;
    return False;

def findProcessesByName(name, processedResults, args = False):
    pids = [];
    for pid in processedResults:
        if match(name, processedResults[pid]['cmd'], args):
            pids.insert(len(pids), pid);
    return pids;

def killProcesses(pids, processedResults, force = False):
    killList = [];
    if type(pids) is type(int()):
        pids = [pids];
    elif type(pids) is type(str()):
        pids = [int(pids)];
    elif type(pids) is type(list()):
        pids = [int(pid) for pid in pids];
    else:
        raise TypeError("killProcesses(): Inappropriate type!");
    
    if not force:
        for pid in pids:
            skip = False;
            for each in warnOn:
                if match(each, processedResults[pid]["cmd"]):
                    print "Warning: Attempt to kill process '" + fix(processedResults[pid]["cmd"], None, 32) + "' (" + str(pid) + ") aborted. Use -f instead.";
                    skip = True;
                    break;
            if not skip:
                killList.insert(len(killList), pid);
    else:
        killList = pids;

    killList = [str(pid) for pid in killList];
    strpids = "";
    strpids = " ".join(killList);
    status, result = commands.getstatusoutput("kill -9 " + strpids);
    return result;

def processResults(results):
    interem = [processLine(line.strip()) for line in result.split("\n")][1:-1];
    results = {int(e[1]):{'uid':e[0], 'ppid':e[2], 'c':e[3], 'stime':e[4], 'tty':e[5], 'time':e[6], 'cmd':e[7]} for e in interem};
    return results;

def fix(string, formatter = "{:<12}", maxlen = 80):
    if formatter is None:
        def formatStr(string):
            return string;
    else:
        def formatStr(string):
            return formatter.format(string);
        
    string = str(string);
    if type(maxlen) is type(int()):
        if len(string) > maxlen:
            string = string[:maxlen - 5];
            string+= "[...]";
    return formatStr(string);

def printProcessed(processedResults, needle = "*", subset = set(["pid", "uid", "cmd"]), formatter = "{:<12}", maxlen = 80, args = False):
    string = "";

    # Alt. Defaults
    ideal = ["uid", "pid", "ppid", "c", "stime", "tty", "time", "cmd"]
    
    if type(subset) is not type(set()):
        if subset is None:
            subset = set(ideal);
        else:
            subset = set(subset);

    subset = list(subset);
    keys = findProcessesByName(needle, processedResults, args);
    keys.sort();

    for each in ideal:
        if each in subset:
            string+= fix(each, formatter, maxlen).upper() + "\t";
    string+="\n" + "-"*len(string) + "\n";
    
    for pid in keys:
        for each in ideal:
            if each in subset:
                if each is "pid":
                    string+= fix(pid, formatter, maxlen) + "\t";
                else:
                    string+= fix(processedResults[pid][each], formatter, maxlen) + "\t";
        string+="\n";
    string+= str(len(keys)) + " results.";
    print string;

def cmdKill(info = False):
    if info:
        string = "Kills processes with names that match [argument].";
        return string;
    global results;
    flagForce = False;
    flagArgs = False;
    if len(sys.argv) > 3:
        # One or more flags present
        if "-f" in sys.argv[1:-1]:
            flagForce = True;
        if "-a" in sys.argv[1:-1]:
            flagArgs = True;

    pids = findProcessesByName(sys.argv[-1], results, flagArgs);
    if len(pids) == 0:
        print "Error: No such process(es) found!";
        return;
    result = killProcesses(pids, results, flagForce);
    if (result is None) or (result == ''):
        print "Successfully killed " + str(len(pids)) + " process(es).";
    else:
        print "Error encountered while attempting to kill " + str(len(pids)) + " process(es):\n" + result;

def cmdList(info = False):
    if info:
        string = "Lists processes with names that match [argument].";
        return string;
    global results;
    global TITLE;

    flagArgs = False;
    flagNoTruncate = False;
    
    if len(sys.argv) > 3:
        # One or more flags present
        if "-a" in sys.argv[1:-1]:
            flagArgs = True;
        if "-nt" in sys.argv[1:-1]:
            flagNoTruncate = True;

    if len(sys.argv) < 3:
        needle = "*";
    else:
        needle = sys.argv[-1];

    if flagNoTruncate:
        maxlen = None;
    else:
        maxlen = 48;

    printProcessed(results, needle, maxlen = maxlen);
    return;
    
def cmdCycle():
    global results;
    flagForce = False;
    return;
    # Not implemented yet

    if len(sys.argv) > 3:
        # One or more flags present
        if "-f" in sys.argv[1:-1]:
            flagForce = True;
        if "-i" in sys.argv[1:-1]:
            flagIgnore = True;

    pids = findProcessesByName(sys.argv[-1], results);
    if len(pids) == 0:
        print "PyPM: No such process(es) found!";
        return;

    redoList = [];
    killResult = killProcesses(pids, results, flagForce);
    #startResult = startProcesses(redoList);

#print "hello, world!";
status, result = commands.getstatusoutput("ps -Af");
status = 0;
#result = file("temp.txt").read();
warnOn = ["/sbin/init", "sshd*", "[*]"];
#print result;
TITLE = "\"PyPM\" - Python Process Manager, (c) 2011 James King (pypm.codeplex.com)\n";

if __name__ == "__main__":
    # Lets get the processed results
    results = processResults(result);
    # Determine what the command is, if any?
    cmds = {'kill':cmdKill, 'list':cmdList};#, 'cycle':cmdCycle};
    showHelp = False;
    
    if len(sys.argv) > 1:
        # Process command
        #string = "PM: " + repr(sys.argv);
        string = "";
        if sys.argv[1] in cmds:
            cmds[sys.argv[1]]();
        else:
            showHelp = True;
    else:
        showHelp = True;
        
    if showHelp:
        if 'idlelib' in sys.modules:
            # Inside IDLE or some other interpreter
            string = TITLE;
            string+= "-"*(len(string) - 1) + "\n";
            string+= "(Interactive Mode)\n";
        else:
            string = TITLE;
            string+= "-"*(len(string) - 1) + "\n";
            string+= "Usage: pm [command] [flags] [argument]\nNote: '?' can be used as a wildcard (like *)\nFlags: -a Include args, -f Force, -nt No Truncate\n\n";
            string+= "Commands:\n";
            for each in cmds:
                string+= fix(each) + "\t" + cmds[each](True) + "\n";
    print string;
